import numpy as np
import scipy.stats as jslib
import torch
import torch.nn.functional as F

from util.seed import set_seed
from defense.get_defense import get_defense
from raw.get_model import get_model

class DefenseEvaluation:
    def __init__(self, defense, suffix = ''):
        self.defense = defense

        # 加载原始模型
        self.raw_model = get_model(defense.model.model_name)

        # 加载测试数据
        test_loader = defense.model.get_test_loader(batch_size = 30, shuffle = False)

        # 加载网络参数
        self.raw_model.load_raw_state()
        defense.load_enhanced_state(suffix = suffix)

        self.raw_model.eval()
        defense.model.eval()

        # 将图像传入原始网络和防御后网络得预测值，并保存真值
        raw_out_list = [] # 原始softmax
        out_list = [] # 防御后softmax
        label_list = [] # 真实标签
        with torch.no_grad():
            for images, labels in test_loader:
                ori_images = images.clone()
                images = defense.transform_input(images) # 对输入图像进行变形

                ori_images = ori_images.to(defense.model.device)
                images = images.to(defense.model.device)
                labels = labels.to(defense.model.device)

                # 原始模型预测
                raw_outs = self.raw_model(ori_images)
                raw_softmax_outs = F.softmax(raw_outs, dim = 1).cpu().numpy()

                # 防御后模型预测
                outs = defense.predict(images)
                softmax_outs = F.softmax(outs, dim = 1).cpu().numpy()

                raw_out_list.extend(raw_softmax_outs)
                out_list.extend(softmax_outs)
                label_list.extend(labels)

        self.raw_outs = np.array(raw_out_list)
        self.outs = np.array(out_list)
        self.labels = np.array(label_list)

    def evaluate(self):
        # compute the classification accuracy of raw model
        correct_prediction_raw = np.equal(np.argmax(self.raw_outs, axis = 1), self.labels)
        acc_raw = np.mean(correct_prediction_raw.astype(float))

        # compute the classification accuracy of defense-enhanced model
        correct_prediction_def = np.equal(np.argmax(self.outs, axis = 1), self.labels)
        acc_def = np.mean(correct_prediction_def.astype(float))

        # compute the Classification Accuracy Variance(CAV)
        cav_result = acc_def - acc_raw

        # find the index of correct predicted examples by defence-enhanced model and raw model
        idx_def = np.squeeze(np.argwhere(correct_prediction_def == True))
        idx_raw = np.squeeze(np.argwhere(correct_prediction_raw == True))
        idx = np.intersect1d(idx_def, idx_raw, assume_unique = True)

        # compute the Classification Rectify Ratio(CRR) & Classification Sacrifice Ratio(CSR)
        num_rectify = len(idx_def) - len(idx)
        crr_result = num_rectify / len(self.labels)

        num_sacrifice = len(idx_raw) - len(idx)
        csr_result = num_sacrifice / len(self.labels)

        if self.defense.ignore_output: # 跳过对网络输出向量的评估
            ccv_result = cos_result = 0
        else:
            # filter the correct prediction results
            pred_def_filter = self.outs[idx]
            pred_raw_filter = self.raw_outs[idx]

            # compute the Classification Confidence Variance(CCV)
            confidence_def = np.max(pred_def_filter, axis = 1)
            confidence_raw = np.max(pred_raw_filter, axis = 1)
            ccv_result = np.mean(np.absolute(confidence_def - confidence_raw)).astype(np.float64)

            # compute the Classification Output Stability(COS)
            M = (pred_def_filter + pred_raw_filter) / 2.
            js_total = 0
            for i in range(len(M)):
                js = 0.5 * jslib.entropy(pred_def_filter[i], M[i]) + 0.5 * jslib.entropy(pred_raw_filter[i], M[i])
                js_total += js
            cos_result = js_total / len(M)

        return acc_raw, acc_def, cav_result, crr_result, csr_result, ccv_result, cos_result

def main(model_name, defense_name, defense_config, suffix = ''):
    # 设置随机数种子
    set_seed()

    # 加载防御
    defense = get_defense(defense_name, model_name, defense_config)

    # 加载评估
    ev = DefenseEvaluation(defense, suffix = suffix)

    # 进行评估
    raw_acc, def_acc, cav, crr, csr, ccv, cos = ev.evaluate()
    print('Acc of Raw Model:\t{:.2f}%'.format(raw_acc * 100))
    print('Acc of Enhanced Model:\t{:.2f}%'.format(def_acc * 100))
    print('CAV:\t{:.2f}%'.format(cav * 100))
    print('CRR:\t{:.2f}%'.format(crr * 100))
    print('CSR:\t{:.2f}%'.format(csr * 100))
    if defense.ignore_output:
        ccv = None
        cos = None
        print('CCV:\tN/A')
        print('COS:\tN/A')
    else:
        print('CCV:\t{:.2f}%'.format(ccv * 100))
        print('COS:\t{:.4f}'.format(cos))

    return {
        'rawAcc': raw_acc,
        'defAcc': def_acc,
        'cav': cav,
        'crr': crr,
        'csr': csr,
        'ccv': ccv,
        'cos': cos
    }
